# analysis/visualization.py

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import networkx as nx
from grakel.kernels import ShortestPath
from grakel import Graph


def display_combined_heatmap(similarity_matrices, category_labels):
    """
    显示所有类别的相似性热图。
    :param similarity_matrices: 相似性矩阵列表
    :param category_labels: 类别标签
    """
    n = len(category_labels)
    combined_matrix = [[0] * n for _ in range(n)]

    idx = 0
    for i in range(n):
        for j in range(i + 1, n):
            if idx < len(similarity_matrices):
                combined_matrix[i][j] = similarity_matrices[idx]
                combined_matrix[j][i] = similarity_matrices[idx]
                idx += 1

    plt.figure(figsize=(10, 8))
    sns.heatmap(combined_matrix, annot=True, cmap='RdBu_r', xticklabels=category_labels, yticklabels=category_labels)
    plt.title('GO Term Similarity Heatmap')
    plt.show()


def draw_network(G, community, ax=None):
    """
    绘制某个社区在图中的结构。
    :param G: NetworkX 图对象
    :param community: 社区节点列表
    :param ax: Matplotlib 的子图对象（可选）
    """
    subG = G.subgraph(community)
    pos = nx.spring_layout(subG)

    if ax is None:
        plt.figure(figsize=(8, 6))
        ax = plt.gca()

    node_colors = ['skyblue' if n in community else 'lightgray' for n in subG.nodes()]
    nx.draw(subG, pos, with_labels=True, node_color=node_colors, edge_color='gray', node_size=200, alpha=0.8, ax=ax)


def plot_community_similarity(communities, file_path="results/fused_network.txt"):
    """
    绘制多个社区的相似性热图。
    """
    G = nx.readwrite.edgelist.read_weighted_edgelist(file_path, nodetype=int)
    fig, axes = plt.subplots(len(communities), len(communities), figsize=(12, 10))

    sp_kernel = ShortestPath(normalize=True, with_labels=False)

    for i, c1 in enumerate(communities):
        for j, c2 in enumerate(communities):
            subG1 = G.subgraph(c1)
            subG2 = G.subgraph(c2)

            adj1 = nx.to_numpy_array(subG1, weight='weight')
            adj2 = nx.to_numpy_array(subG2, weight='weight')

            g1 = Graph(adj1)
            g2 = Graph(adj2)
            sim = sp_kernel.fit_transform([g1])
            sim = sp_kernel.transform([g2])

            axes[i][j].imshow([[sim[0][0]]], cmap='coolwarm', interpolation='none')
            axes[i][j].set_title(f"C{i + 1} vs C{j + 1}")
            axes[i][j].axis('off')

    plt.tight_layout()
    plt.savefig("results/figures/community_similarity.png")
    plt.show()


# analysis/visualization.py

def plot_community_similarity_matrix(communities, output_file="results/figures/community_similarity_matrix.png", G=None):
    """
    绘制所有社区之间的相似性矩阵热图。
    :param communities: 社区列表 [[nodes]]
    :param output_file: 输出路径
    """
    n = len(communities)
    sim_matrix = np.zeros((n, n))

    sp_kernel = ShortestPath(normalize=True, with_labels=False)

    graphs = [Graph(nx.to_numpy_array(nx.subgraph(G, c))) for c in communities]
    sp_kernel.fit_transform(graphs)

    for i in range(n):
        for j in range(i + 1, n):
            sim = sp_kernel.transform([graphs[j]])
            sim_matrix[i][j] = sim[0][0]
            sim_matrix[j][i] = sim[0][0]

    plt.figure(figsize=(8, 6))
    sns.heatmap(sim_matrix, annot=True, cmap='coolwarm')
    plt.title("社区相似性矩阵")
    plt.savefig(output_file, dpi=300)
    plt.close()
